import pandas as pd
import numpy as np
import h5py
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.python.framework import ops
import time
from tensorflow import keras
from tensorflow.keras.utils import to_categorical
import os
from network_function import mymodel, get_data,plot_confuse

# -----------------参数设置
# datatype='dist_feature'
datatype = 'all_features'
checkpointpath = 'D:/work_dxx/holistic/best_model_fconv_5_' + datatype + '.weights'
# checkpointpath='best_model_fconv_all_features.weights'
modelpath = 'D:/work_dxx/holistic/mymodel_fconv_5_' + datatype +'.h5'
bestmodelpath = 'D:/work_dxx/holistic/bestmodel_fconv_5_' + datatype +'.h5'

# ---------getdata-------------
train_data, test_data = get_data(datatype)
# train_data
x_train = np.array(train_data)[:, 1:]
y_train = np.array(train_data)[:, 0]
y_train = to_categorical(y_train, 7)
# test_data
x_test = np.array(test_data)[:, 1:]
y_test_orig = np.array(test_data)[:, 0]
y_test = to_categorical(y_test_orig, 7)
# 输入data维度
x_train = np.expand_dims(x_train, axis=2) #表示是是增加的维度是在第三个维度上
x_test = np.expand_dims(x_test, axis=2) #表示是是增加的维度是在第三个维度上

input_shape = (x_train.shape[1],1)

print(f'y_train.shape={y_train.shape}')  # y_train.shape=(7, 53)
print(f'x_train.shape={x_train.shape}')  # x_train.shape=(9, 53)
print(f'y_test.shape={y_test.shape}')  # y_test.shape=(7, 17)
print(f'x_test.shape={x_test.shape}')  # x_test.shape=(9, 17)


# 开始训练
def train(epochs=100, dropout1=0.1, dropout2=0.3):
    print('building model------------------')
    model, history = mymodel(input_shape)
    print('Training ------------')
    # 保存最好的模型
    checkpoint = keras.callbacks.ModelCheckpoint(checkpointpath,
                                                 monitor='val_acc',
                                                 verbose=1,
                                                 save_best_only=True,
                                                 mode='auto',
                                                 period=1)
    # model.load_weights(checkpointpath)
    start_time = time.perf_counter()
    # model.fit(x_train, y_train, epochs=epochs, batch_size=32,callbacks=[history], validation_split=0.2)
    model.fit(x_train, y_train, epochs=epochs, batch_size=32, callbacks=[history, checkpoint], validation_data=(x_test, y_test))
    end_time = time.perf_counter()
    # 计算时差
    print("CPU的执行时间 = " + str(end_time - start_time) + " 秒")
    tf.keras.models.save_model(model, modelpath, save_format='h5')
    print('模型已保存')
    # 绘制acc-loss曲线
    history.loss_plot('epoch')


# 评价训练出的网络
def test(load_model_path):
    print('\nTesting ------------')
    # reconstructed_model = keras.models.load_model("mymodel_all_features")
    reconstructed_model = keras.models.load_model(modelpath)
    # reconstructed_model.load_weights("best_model_conv_all_features.weights")
    # reconstructed_model.load_weights(checkpointpath)
    try:
        tf.keras.models.save_model(reconstructed_model, bestmodelpath, save_format='h5')
        print(f'best model saved in {bestmodelpath}')
    except:
        pass
    # y_test_hat = np.argmax(reconstructed_model.predict(x_test), axis=1)
    # print(y_test_hat.shape)
    # print(y_test_orig.astype(np.int), '\n', y_test_hat)
    # sess=tf.Session()
    # print(tf.argmax(y_test_hat,axis=1).eval(session=sess),'\n',y_test)
    loss, accuracy = reconstructed_model.evaluate(x_test, y_test)
    print('test loss: ', loss)
    print('test accuracy: ', accuracy)
    # plot_confuse(reconstructed_model, x_test, y_test)
    plot_confuse(reconstructed_model, x_train, y_train,figname='train_plot')
    plot_confuse(reconstructed_model, x_test, y_test,figname='test_plot')


mode = 'train'
# mode='test'
if __name__ == '__main__':
    if mode == 'train':
        train(epochs=300, dropout1=0.1, dropout2=0.2)
        test(modelpath)

    elif mode == 'test':
        test(bestmodelpath)